#pragma once

#include <iostream>
#include <assert.h>
#include <time.h>
using namespace std;

enum Colour // 枚举颜色
{
	RED,
	BLACK
};

template<class T>
struct RBTreeNode
{
public:
	RBTreeNode<T>* _left;
	RBTreeNode<T>* _right;
	RBTreeNode<T>* _parent;

	T _data; // 结点中的数据
	Colour _col; // 比AVL树少了平衡因子，多了颜色

	RBTreeNode(const T& data)
		:_left(nullptr)
		, _right(nullptr)
		, _parent(nullptr)
		, _data(data)
	{}
};

template<class T, class Ref, class Ptr>
struct __RBTreeIterator
{
public:
	typedef RBTreeNode<T> Node;
	typedef __RBTreeIterator<T, Ref, Ptr> Self;
	Node* _node;

	__RBTreeIterator(Node* node)
		:_node(node)
	{}

	Ref operator*()
	{
		return _node->_data;
	}

	Ptr operator->()
	{
		return &_node->_data;
	}

	bool operator!=(const Self& s) const
	{
		return _node != s._node;
	}

	bool operator==(const Self& s) const
	{
		return _node == s._node;
	}

	Self& operator++()
	{
		if (_node->_right) // 右子树存在，++后就到右子树的最左结点
		{
			Node* left = _node->_right;
			while (left->_left)
			{
				left = left->_left;
			}

			_node = left;
		}
		else // 右子树不存在，++后就到：所在子树是左子树的最近祖宗节点。
		{
			Node* parent = _node->_parent; // 找祖先里面孩子不是祖先的右的那个
			Node* cur = _node;
			while (parent && cur == parent->_right)
			{   // 是右子树就往上走，parent存在是处理_node是中序最后一个结点的情况
				cur = cur->_parent;
				parent = parent->_parent;
			}

			_node = parent;
		}

		return *this;
	}

	Self operator++(int)
	{
		Self ret = Self(_node); // 记录当前位置 最后返回

		if (_node->_right) // 右子树存在，++后就到右子树的最左结点
		{
			Node* left = _node->_right;
			while (left->_left)
			{
				left = left->_left;
			}

			_node = left;
		}
		else // 右子树不存在，++后就到：所在子树是左子树的最近祖宗节点。
		{
			Node* parent = _node->_parent; // 找祖先里面孩子不是祖先的右的那个
			Node* cur = _node;
			while (parent && cur == parent->_right)
			{   // 是右子树就往上走，parent存在是处理_node是中序最后一个结点的情况
				cur = cur->_parent;
				parent = parent->_parent;
			}

			_node = parent;
		}

		return ret;
	}

	Self& operator--()
	{
		if (_node->_left) // 左子树存在，++后就到左子树的最右结点
		{
			Node* right = _node->_left;
			while (right->_right)
			{
				right = right->_right;
			}

			_node = right;
		}
		else // 左子树不存在，++后就到：所在子树是右子树的最近祖宗节点。
		{
			Node* parent = _node->_parent; // 找祖先里面孩子不是祖先的左的那个
			Node* cur = _node;
			while (parent && cur == parent->_left)
			{  // 是左子树就往上走，parent存在是处理_node是中序第一个结点的情况
				cur = cur->_parent;
				parent = parent->_parent;
			}

			_node = parent;
		}

		return *this;
	}

	Self operator--(int)
	{
		Self ret = Self(_node); // 记录当前位置 最后返回

		if (_node->_left) // 左子树存在，++后就到左子树的最右结点
		{
			Node* right = _node->_left;
			while (right->_right)
			{
				right = right->_right;
			}

			_node = right;
		}
		else // 左子树不存在，++后就到：所在子树是右子树的最近祖宗节点。
		{
			Node* parent = _node->_parent; // 找祖先里面孩子不是祖先的左的那个
			Node* cur = _node;
			while (parent && cur == parent->_left)
			{  // 是左子树就往上走，parent存在是处理_node是中序第一个结点的情况
				cur = cur->_parent;
				parent = parent->_parent;
			}

			_node = parent;
		}

		return ret;
	}
};

template<class K, class T, class KeyOfT> // KeyOfT仿函数，把key取出来
struct RBTree
{
	typedef RBTreeNode<T> Node;
public:
	typedef __RBTreeIterator<T, T&, T*> iterator;

	iterator begin()
	{
		Node* left = _root;
		while (left && left->_left)
		{
			left = left->_left;
		}

		return iterator(left);
	}

	iterator end()
	{
		return iterator(nullptr);
	}

	pair<iterator, bool> Insert(const T& data) // 1. 返回值改成pair<iterator, bool>
	{
		KeyOfT kot; // 定义一个仿函数的对象

		if (_root == nullptr)
		{
			_root = new Node(data);
			_root->_col = BLACK; // 根给黑色
			return make_pair(iterator(_root), true); // 2. 空树时返回根的迭代器和true的键值对
		}

		Node* parent = nullptr;
		Node* cur = _root;
		while (cur) // 找到要插入的结点
		{
			if (kot(cur->_data) < kot(data))
			{
				parent = cur;
				cur = cur->_right;
			}
			else if (kot(cur->_data) > kot(data))
			{
				parent = cur;
				cur = cur->_left;
			}
			else
			{
				return make_pair(cur, false); // 3. 存在新插入的数据，返回原本存在的数据的迭代器和false
			}
		}

		cur = new Node(data);
		Node* newnode = cur; // 记录最后插入成功返回新插入结点的迭代器
		cur->_col = RED; // 默认插入红色结点
		if (kot(parent->_data) < kot(data)) // 找到位置后插入结点
		{
			parent->_right = cur;
		}
		else
		{
			parent->_left = cur;
		}
		cur->_parent = parent;

		while (parent && parent->_col == RED) // 父亲存在且为红才需要处理
		{
			Node* grandfather = parent->_parent;
			assert(grandfather); // 确定的可以断言下，否则就是插入前就不是红黑树
			assert(grandfather->_col == BLACK);
			if (grandfather->_left == parent)
			{
				Node* uncle = grandfather->_right;

				if (uncle && uncle->_col == RED) // 情况一，叔叔存在且为红(可以直接复制到下面uncle在左边)
				{    // 将父亲和叔叔改为黑，祖父改为红，然后把祖父当成cur，parent变祖父parent继续向上调整。
					parent->_col = uncle->_col = BLACK;
					grandfather->_col = RED;
					cur = grandfather;
					parent = cur->_parent;
				}
				else // 情况二或情况三：叔叔存在且为黑或叔叔不存在
				{
					if (cur == parent->_left) // 情况二的右旋+变色(parent在左)
					{
						//     g      
						//   p   u
						// c
						RotateR(grandfather);
						parent->_col = BLACK; // 父亲变为根了
						grandfather->_col = RED;
					}
					else // 情况二的左右双旋+变色(parent在左)
					{
						//      g      
						//   p     u
						//    c
						RotateL(parent);
						RotateR(grandfather);
						cur->_col = BLACK; // cur变为根了
						grandfather->_col = RED;
					}
					break;
				}
			}
			else
			{
				Node* uncle = grandfather->_left;

				if (uncle && uncle->_col == RED) // 情况一，叔叔存在且为红
				{    // 将父亲和叔叔改为黑，祖父改为红，然后把祖父当成cur，parent变祖父parent继续向上调整。
					parent->_col = uncle->_col = BLACK;
					grandfather->_col = RED;
					cur = grandfather;
					parent = cur->_parent;
				}
				else // 情况二或情况三：叔叔存在且为黑或叔叔不存在
				{
					if (cur == parent->_right) // 情况二的左旋+变色(parent在右)
					{
						//     g      
						//   u   p
						//        c
						RotateL(grandfather);
						parent->_col = BLACK; // 父亲变为根了
						grandfather->_col = RED;
					}
					else // 情况二的右左双旋+变色(parent在右)
					{
						//       g      
						//    u     p
						//         c
						RotateR(parent);
						RotateL(grandfather);
						cur->_col = BLACK; // cur变为根了
						grandfather->_col = RED;
					}
					break;
				}
			}
		}

		_root->_col = BLACK;
		return make_pair(newnode, true); // 4. 插入成功，返回新插入数据的迭代器和true
	}

	void InOrder()
	{
		_InOrder(_root);
		cout << endl;
	}

	bool IsBalance()
	{
		if (_root == nullptr)
		{
			return true;
		}

		if (_root->_col == RED) // 验证性质二
		{
			cout << "根节点不是黑色" << endl;
			return false;
		}

		int benchmark = 0; // 黑色节点数量基准值
		//Node* cur = _root; // 这种方法是先遍历一遍，然后传值，不过我们可以传引用
		//while (cur)
		//{
		//	if (cur->_col == BLACK)
		//	{
		//		++benchmark;
		//	}
		//	cur = cur->_left;
		//}
		return PrevCheck(_root, 0, benchmark); // 验证性质三和四
	}

protected:
	bool PrevCheck(Node* root, int blackNum, int& benchmark)
	{
		if (root == nullptr)
		{
			if (benchmark == 0)
			{
				benchmark = blackNum;
				return true;
			}

			if (blackNum != benchmark) // 验证性质三
			{
				cout << "某条黑色节点的数量不相等" << endl;
				return false;
			}
			else
			{
				return true;
			}
		}

		if (root->_col == BLACK)
		{
			++blackNum;
		}

		if (root->_col == RED && root->_parent->_col == RED) // 验证性质四
		{
			cout << "存在连续的红色节点" << endl;
			return false;
		}

		return PrevCheck(root->_left, blackNum, benchmark)
			&& PrevCheck(root->_right, blackNum, benchmark);
	}

	void _InOrder(Node* root)
	{
		if (root == nullptr)
		{
			return;
		}
		_InOrder(root->_left);
		cout << root->_kv.first << ":" << root->_kv.second << endl;
		_InOrder(root->_right);
	}

	void RotateL(Node* parent)
	{
		Node* subR = parent->_right; // 动了三个标记了的结点，共更新六个指针，这更新两个指针
		Node* subRL = subR->_left;

		parent->_right = subRL;
		if (subRL) // subRL不为空才更新
		{
			subRL->_parent = parent;
		}

		Node* ppNode = parent->_parent; // 记录parent的parent，防止parent是一颗子树的头结点

		subR->_left = parent; // 再更新两个指针
		parent->_parent = subR;

		if (_root == parent)  // 最后更新两个指针
		{
			_root = subR;
			subR->_parent = nullptr;
		}
		else // parent是一颗子树的头结点
		{
			if (ppNode->_left == parent)
			{
				ppNode->_left = subR;
			}
			else
			{
				ppNode->_right = subR;
			}
			subR->_parent = ppNode;
		}
	}

	void RotateR(Node* parent)
	{
		Node* subL = parent->_left;
		Node* subLR = subL->_right;

		parent->_left = subLR; // 更新两个节点
		if (subLR)
		{
			subLR->_parent = parent;
		}

		Node* ppNode = parent->_parent;

		subL->_right = parent; // 再更新两个节点
		parent->_parent = subL;

		if (_root == parent) // 最后更新两个结点
		{
			_root = subL;
			subL->_parent = nullptr;
		}
		else
		{
			if (ppNode->_left == parent)
			{
				ppNode->_left = subL;
			}
			else
			{
				ppNode->_right = subL;
			}

			subL->_parent = ppNode;
		}
	}

	Node* _root = nullptr;
};